'''
单张数据出图
'''
from matplotlib import pyplot as plt
import matplotlib
import os 
import numpy as np
import sys
'''torch'''
import torch
import torch.nn as torch_nn
import torch.nn.functional as F
import torchvision.transforms as transforms
'''ourfunction'''
from models.get_models import get_models
from utils.utils import load_ckpt,getedge,getedge_sig
from datasetloader.get_datasetLoader import get_augmentations
from datasetloader.augmentations import BandVI,OneHot,Normalize,ToTensor

'''
使用matplotlib 作为图片绘制工具
整体结构，方法组

- drawimage 直接绘制图片
- predictOne 预测单张图片
- loadmodel 加载模型
- pre_augmentation  预处理模块
- main_process 主要调用方法
'''

def LoadModel(modelName,modelcpk_path,num_cls,input_channels):
    '''
    加载模型
    args:
        modelName 需要参与计算的模型的名称
        modelcpk_path 训练好模型的保存路径
    return:
          model : 加载好对应权重的模型
    '''
    model=get_models(modelName,num_class=num_cls,channels_num=input_channels)
    epoch,model,best_iou=load_ckpt(model,None,modelcpk_path,'cuda')
    return epoch,model,best_iou

def pre_augmentation(cls_num,mean,std,VI_enable=False,minvalue=None,maxvalue=None):
    if VI_enable:
        pre_val_list=[ 
        BandVI(enable=VI_enable,minvalue=minvalue,maxvalue=maxvalue),
        Normalize(mean,std,VI_enable=VI_enable) ]
    else:
        pre_val_list=[ 
        Normalize(mean,std,VI_enable=VI_enable)
        ]
    return transforms.Compose(pre_val_list)

def predictOne(img,model):
    '''
    预测单张图像的结果
    args:  
        img narray [H,W,C] 原始图像
        model 模型 
    '''
    # 数据转换
    h,w,c=img.shape
    img_tranpose=np.transpose(img,(2,0,1))
    img_tensor=torch.from_numpy(img_tranpose)
    #img_tensor=torch.transpose(img_tensor,(2,0,1))
    img_tensor=img_tensor.reshape(1,c,h,w)
    img_tensor=img_tensor.cuda().float()
    # 模型预测
    img_pred=model(img_tensor)
    # 概率函数计算
    img_pred=F.softmax(img_pred)
    img_mask=img_pred.max(1)[1]
    return img_mask.detach().cpu().squeeze().numpy()

def drawimage(img,Normal_img,label,pred_mask,savePath):
    '''
    绘制图像
    args:
        img,narray,float,[H,W,C] 原始图像值，其中值的范围为 [0,1]
        label,narray,int [H,W],标签值，数据应不能超过255
        pred_mask,narray,[H,W]，预测值，其中值的范围应与label一致
        savepath,str， 图像保存地址
    return：
        返回对应的地址值
    '''
    # 获得对应的边界结果
    seg_pred=getedge_sig(pred_mask) # 预测的边界情况
    seg=getedge_sig(label) # 真实的边界情况
    # 自定义matplotlib 绘图色带
    
    '''
    绘制图像，3行3列

    【        img,label    ,seg  #   原始模型
              img,pred_mask,pred_seg # 预测结果
              img,label_pred,pred_seg_seg # 叠加分析结果
    】
    '''
    # 构建对比结果
    # 预测结果对比
    h,w=label.shape
    mask_pred_source=np.zeros((h,w,3))
    mask_pred_source[:,:,0]=label
    mask_pred_source[:,:,1]=pred_mask
    # 边界结果对比
    h,w=seg.shape
    seg_pred_source=np.zeros((h,w,3))
    seg_pred_source[:,:,0]=seg
    seg_pred_source[:,:,1]=seg_pred
    # 处理原始mask与seg
    label=np.copy(mask_pred_source)
    label[:,:,1]=label[:,:,1]*0
    seg=np.copy(seg_pred_source)
    seg[:,:,1]=seg[:,:,1]*0
    # 处理预测影像
    pred_mask=np.copy(mask_pred_source)
    pred_mask[:,:,0]=pred_mask[:,:,0]*0
    seg_pred=np.copy(seg_pred_source)
    seg_pred[:,:,0]=seg_pred[:,:,0]*0

    
    fig, ax = plt.subplots(3, 3, figsize=(9, 9),dpi=300)
    # 展示原图像
    ax[0,0].imshow(img)
    ax[0,1].imshow(label)
    ax[0,2].imshow(seg)
    # 展示结果图像
    ax[1,0].imshow(Normal_img.astype(np.uint8))
    ax[1,1].imshow(pred_mask)
    ax[1,2].imshow(seg_pred)
    # 展示对比结果
    ax[2,0].imshow(img)
    ax[2,1].imshow(mask_pred_source) 
    ax[2,2].imshow(seg_pred_source)
    for a in ax.ravel():
        a.set_axis_off() 
    plt.subplots_adjust( wspace=0, hspace=0)
    plt.savefig(savePath,dpi=300)    
    plt.show()
      
    
def main_process(config_param):
    '''
    处理图像的主要流程文件
    '''
    img=np.load(config_param['img'])
    label=np.load(config_param['label'])
    augmentation=pre_augmentation(config_param['cls_num'],
                                    config_param['mean'],config_param['std'],
                                    VI_enable=config_param['VI_enable'],
                                    minvalue=config_param['minvalue'],
                                    maxvalue=config_param['maxvalue'])
    epoch,model,best_iou=LoadModel(config_param['modelname'],config_param['checkpoint'],config_param['cls_num'],config_param['chns_num'])
    print("epoch:{} , best_iou:{}".format(epoch,best_iou))
    # 调整模型
    model=model.cuda()
    model.eval()
    with torch.no_grad():
        # 预测模型
        sample={'img':img,"seg":label,"label":label,"name":"",'source_image':np.copy(img)}
        sample=augmentation(sample)
        Normal_img=sample['img']
        pred_mask=predictOne(Normal_img,model)
        # 绘制图形
        min_norm=np.min(np.min(Normal_img))
        max_norm=np.max(np.max(Normal_img))
        Normal_img=255*(Normal_img-min_norm)/(max_norm-min_norm)
        Normal_img=Normal_img.astype(np.uint8)
        drawimage(img,Normal_img,label,pred_mask,config_param['savePath'])



if __name__=='__main__':
    config_param={
        'img':'/media/gis/databackup/ayc/modellist/dataset/nanchang/E512/RGB123/test/img/1616.npy',
        'label':'/media/gis/databackup/ayc/modellist/dataset/nanchang/E512/RGB123/test/label/1616.npy',
        'cls_num':2,
        'mean':[98.42966601004996, 103.47446405573739, 89.24539776311728],
        'std':[39.68257615485789, 40.96571881718136, 36.232309241791654],
        'VI_enable':False,
        'minvalue':-1262.0, # 指数运算时使用
        'maxvalue':9483.0,
        'chns_num':3,
        'modelname':'fcn',
        'checkpoint':'/media/gis/databackup/ayc/model_hastrained_log/RGB123_fcn_result/checkpoint/Bestest_mIou_ckpt_fcn_dataName_RGB123.pth',
        'savePath':'/media/gis/databackup/ayc/modellist/ResultImg/fcn_RGB123.png',
    }
    main_process(config_param)
 

